from utils.interval import Interval
from utils.misc import INF

from opt.solver import Solver


class MCTS_ParamSet(Solver.ParamSet):
    node_factory = Solver.Param(description="function that creates the root node",
                                options="callable returning a TreeNode",
                                default=None)

    @node_factory.adapter
    def node_factory(self, fnc):
        if fnc is None:
            return type(self).node_factory
        if callable(fnc):
            return fnc
        raise Exception("node factory should be callable or None")

    # --------------------------------------------------------------------------
    metanode_factory = Solver.Param(description="function used to create metanodes",
                                    options="callable returning a MetaNode",
                                    default=None)

    @metanode_factory.adapter
    def metanode_factory(self, fnc):
        if fnc is None:
            return type(self).metanode_factory
        if callable(fnc):
            return fnc
        raise Exception("metanode factory should be callable or None")

    # --------------------------------------------------------------------------
    exploration_coeff = Solver.Param(description=("controls balance between exploration"
                                                  " and exploitation in UCT"),
                                     domain=Interval(0.0, INF),
                                     default=1.0)

    @exploration_coeff.adapter
    def exploration_coeff(self, coeff):
        return float(coeff)

    # --------------------------------------------------------------------------
    expansion_limit = Solver.Param(description="the maximum of children created per iteration",
                                   domain=Interval(1, INF),
                                   default=1)

    @expansion_limit.adapter
    def expansion_limit(self, limit):
        return float(limit)

    # --------------------------------------------------------------------------
    selection_score_terms = Solver.Param(description=("list of terms used in the node "
                                                      "selection score expression"),
                                         default="exploitation_scores exploration_scores_plain")

    @selection_score_terms.adapter
    def selection_score_terms(self, terms):
        """We can pass a list of callables, strings, or a single string."""
        if isinstance(terms, basestring):
            terms = terms.split()
        solver = self.object
        actual_terms = []
        for term in terms:
            if isinstance(term, str):
                actual_terms.append(getattr(solver.MetaNode, term))
            elif callable(term):
                actual_terms.append(term)
            else:
                raise ValueError("score terms must be callable or string: {!r}".format(term))
        return actual_terms

    # --------------------------------------------------------------------------
    exploit_score_fnc = Solver.Param(description="function used to compute exploitation scores",
                                     options="exponential, linear",
                                     default="linear")

    @exploit_score_fnc.adapter
    def exploitation_score_fnc(self, fnc):
        solver = self.object
        if isinstance(fnc, basestring):
            return getattr(solver.MetaNode, "exploit_score_" + fnc)
        elif callable(fnc):
            return fnc
        else:
            raise TypeError("expected callable or string: {!r}".format(fnc))
